package sm4

import (
	"errors"

	"codeup.aliyun.com/lsxd/zltx/adapter/utils/cryptx/cipher"
	"codeup.aliyun.com/lsxd/zltx/adapter/utils/formater"
	"github.com/emmansun/gmsm/sm4"
)

type Option func(opt *options)

// WithFormat 设置输出格式
func WithFormat(f formater.Formater) Option {
	return func(opt *options) {
		opt.format = f
	}
}

// WithCipher 设置加密器
func WithCipher(c *cipher.Cipher) Option {
	return func(opt *options) {
		opt.cipher = c
	}
}

type options struct {
	format formater.Formater
	cipher *cipher.Cipher
}

type SM4 struct {
	o *options
}

// NewSm4 创建一个新的SM4加密器
func NewSm4(opts ...Option) *SM4 {
	opt := &options{
		format: formater.Base64,
		cipher: cipher.NewCipher(),
	}
	for _, o := range opts {
		o(opt)
	}
	return &SM4{o: opt}
}

// Encrypt 加密数据
func (d *SM4) Encrypt(data string) (string, error) {
	if len(data) == 0 {
		return "", nil
	}
	content := []byte(data)
	key := d.o.cipher.GetKey(0)
	iv := d.o.cipher.GetKey(1)
	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}
	if d.o.cipher.Mode != cipher.ECB && len(iv) != block.BlockSize() {
		return "", errors.New("invalid iv length")
	}
	if d.o.cipher.Padding == cipher.No && len(content)%block.BlockSize() != 0 {
		return "", errors.New("invalid padding")
	}

	dst, err := d.o.cipher.Cipher().Encrypt(content, block)
	if err != nil {
		return "", err
	}
	return d.o.format.Encode(dst), nil
}

// Decrypt 解密数据
func (d *SM4) Decrypt(data string) (string, error) {
	content, err := d.o.format.Decode(data)
	if err != nil {
		return "", err
	}
	if len(content) == 0 {
		return "", nil
	}

	key := d.o.cipher.GetKey(0)
	iv := d.o.cipher.GetKey(1)
	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}
	if d.o.cipher.Mode != cipher.ECB && len(iv) != block.BlockSize() {
		return "", errors.New("invalid iv length")
	}
	if (d.o.cipher.Mode == cipher.CBC || d.o.cipher.Padding == cipher.No) && len(content)%block.BlockSize() != 0 {
		return "", errors.New("invalid padding")
	}

	dst, err := d.o.cipher.Cipher().Decrypt(content, block)
	if err != nil {
		return "", err
	}
	return string(dst), nil
}
